package javaagent;

import javassist.*;

import java.io.IOException;
import java.lang.instrument.*;
import java.security.ProtectionDomain;
import java.util.Arrays;

/**
 * @author ：zzs
 * @version : 1.0
 * @date ：Created in 2022/2/24 19:00
 * @description：启动类(加载启动)
 */
public class MyAgent {
    /**
     * @param arg             启动参数
     * @param instrumentation 类的加载和定义都是通过它来实现的
     */
    // 加载时启动
    public static void premain(String arg, Instrumentation instrumentation) throws UnmodifiableClassException, ClassNotFoundException {
        System.out.println("hello premain");
        // 如果类在变压器之前加载那么就不能修改类
        HelloWorld helloWorld = new HelloWorld();
        instrumentation.addTransformer(new ClassFileTransformer() {
            /**
             * @param loader
             * @param className           动态加载的类名
             * @param classBeingRedefined 这个类重新加载之前的类
             * @param protectionDomain    类的基本信息
             * @param classfileBuffer     这个类的字节码，如果返回null就按照原有的进行加载覆盖
             * @return 返回指令码
             * @throws IllegalClassFormatException
             */
            @Override
            public byte[] transform(ClassLoader loader, String className,
                                    Class<?> classBeingRedefined,
                                    ProtectionDomain protectionDomain,
                                    byte[] classfileBuffer) throws IllegalClassFormatException {
                if (!"javaagent/HelloWorld".equals(className)) {
                    return null;
                }
                // 配合javassist修改指令码
                ClassPool pool = new ClassPool();
                pool.appendSystemPath();
                try {
                    CtClass ctClass = pool.get("javaagent.HelloWorld");
                    CtMethod method = ctClass.getDeclaredMethod("hello");
                    method.insertBefore("System.out.println(\"插入前置逻辑\");");
                    return ctClass.toBytecode();
                } catch (NotFoundException | CannotCompileException | IOException e) {
                    e.printStackTrace();
                }
                return null;
            }
        }, true);
        System.out.println(Arrays.toString(instrumentation.getAllLoadedClasses()));
        // 重新触发指定类的加载但是对字节码的修改有一定的限制，只是重走变压器的逻辑
        instrumentation.retransformClasses(HelloWorld.class);
        // 这里的访问的方法内部逻辑已经被修改了
        // 原理有点像是spring热部署:(单例)类只加载一次它所指向的方法地址不变，但是方法自身指令码发生了改变（钥匙还是原来的钥匙但是房间里的东西可能就不一样了！）！
        helloWorld.hello();

        /**
         * 重新定义示例
         */
//        instrumentation.addTransformer(new ClassFileTransformer() {
//            @Override
//            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
//                // 返回的是重新定义后的类
//                return null;
//            }
//        });
//        try {
//            ClassPool pool = new ClassPool();
//            pool.appendSystemPath();
//            CtClass ctClass = pool.get("javaagent.HelloWorld");
//            CtMethod method = ctClass.getDeclaredMethod("hello");
//            method.insertAfter("System.out.println(\"插入后置逻辑\");");
//            instrumentation.redefineClasses(new ClassDefinition(HelloWorld.class, ctClass.toBytecode()));
//        } catch (CannotCompileException | NotFoundException | IOException e) {
//            e.printStackTrace();
//        }
    }


    // 运行时启动
    public static void agentmain(String arg, Instrumentation instrumentation) {
        System.out.println("hello agentmain22");
        HelloWorld helloWorld = new HelloWorld();
        instrumentation.addTransformer(new ClassFileTransformer() {
            /**
             * @param loader
             * @param className           动态加载的类名
             * @param classBeingRedefined 这个类重新加载之前的类
             * @param protectionDomain    类的基本信息
             * @param classfileBuffer     这个类的字节码，如果返回null就按照原有的进行加载覆盖
             * @return 返回指令码
             * @throws IllegalClassFormatException
             */
            @Override
            public byte[] transform(ClassLoader loader, String className,
                                    Class<?> classBeingRedefined,
                                    ProtectionDomain protectionDomain,
                                    byte[] classfileBuffer) throws IllegalClassFormatException {
                if (!"javaagent/HelloWorld".equals(className)) {
                    return null;
                }
                // 配合javassist修改指令码
                ClassPool pool = new ClassPool();
                pool.appendSystemPath();
                try {
                    CtClass ctClass = pool.get("javaagent.HelloWorld");
                    CtMethod method = ctClass.getDeclaredMethod("hello");
                    method.insertBefore("System.out.println(\"插入前置逻辑2\");");
                    return ctClass.toBytecode();
                } catch (NotFoundException | CannotCompileException | IOException e) {
                    e.printStackTrace();
                }
                return null;
            }
        }, true);
        helloWorld.hello();
    }
}
